{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import pandas as pd\n",
    "from sqlalchemy import create_engine\n",
    "import ipywidgets as widgets\n",
    "from IPython.display import display\n",
    "from IPython.display import update_display\n",
    "from IPython.display import display_pretty\n",
    "\n",
    "\n",
    "SQLALCHEMY_DATABASE_URI = os.environ.get('SQLALCHEMY_DATABASE_URI')\n",
    "engine = create_engine(SQLALCHEMY_DATABASE_URI)\n",
    "\n",
    "sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))\n",
    "from api import GPT, Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Define display flow for GPT question input\n",
    "def on_button_clicked(b):\n",
    "    def sql_button_clicked(b):\n",
    "        df = pd.read_sql(query, engine)\n",
    "        df_pretty = df.style.hide_index().set_properties(**{'background-color': 'black',\n",
    "                               'color': 'lawngreen',\n",
    "                               'border-color': 'white'})\n",
    "        display(df_pretty)\n",
    "        \n",
    "    print ('\\033[1mInput:\\033[0m ' + inp.value)\n",
    "    output = gpt.submit_request(inp.value)\n",
    "    result = output['choices'][0].text\n",
    "    query = result.split('output:')[1]\n",
    "    print ('\\033[1mGPT-3 Response:\\033[0m ' + query)\n",
    "    button2 = widgets.Button(description=\"Run SQL\")\n",
    "    button2.on_click(sql_button_clicked)\n",
    "    display(button2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Construct GPT-3-instruct instance, add instruction and examples\n",
    "gpt = GPT(engine=\"davinci-instruct-beta\",\n",
    "          temperature=0.3,\n",
    "          max_tokens=200)\n",
    "gpt.add_instruction('Given an input question, respond with syntactically correct PostgreSQL.')\n",
    "\n",
    "gpt.add_example(Example('select columns from users table', \n",
    "                        'select id, email, dt, plan_type, created_on, updated_on from users'))\n",
    "gpt.add_example(Example('select columns from the charges table', \n",
    "                        'select amount, amount_refunded, created, customer_id, status from charges'))\n",
    "gpt.add_example(Example('select columns from the customers table', \n",
    "                        'select created, email, id from customers'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1mInstruction:\u001b[0m Given an input question, respond with syntactically correct PostgreSQL.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7fe03e81968e4d08b8a63ff74f7fc777",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(Text(value='', description='Ask GPT-3:'), Button(description='Get GPT-3 Response', style=Button…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#Display UI to give GPT-3 prompt and run and display resulting SQL\n",
    "inp = widgets.Text(description='Ask GPT-3:')\n",
    "button1 = widgets.Button(description=\"Get GPT-3 Response\")\n",
    "Box = widgets.HBox([inp, button1])\n",
    "print ('\\033[1mInstruction:\\033[0m ' + gpt.get_instruction_text())\n",
    "button1.on_click(on_button_clicked)\n",
    "display(Box)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
